# -*- coding: utf-8 -*-  
'''
前置处理相关函数

@author: luoyi
Created on 2021年6月23日
'''
import numpy as np

import utils.conf as conf
import utils.relationships as rel


#    实体头尾握手信息转换为相对∑max_sen_len的向量索引信息
def et2et_vidx(text, postag, max_sen_len=conf.TPLINKER.get_max_sentence_len()):
    '''将词性标注信息转换为相对∑max_sen_len的向量索引信息
        @param text: 原文本字符串
        @param postag: [{word:实体词字符串, loc:(头尾指针), pos:词性}...]
        @param max_sen_len: 最大句子长度
        @return idx_et2et: narray(握手指针数量) 实体头尾指针相对于∑max_sen_len向量的索引。loc_s * max_sen_len + loc_e
                idx_no: narray(握手指针数量) 没有et2et握手的索引
    '''
    #    初始化len * len矩阵
    real_len = len(text)
    m = np.zeros([real_len, real_len])
    #    扩充为max_sen_len * max_sen_len矩阵
    m = np.pad(m, [[0, max_sen_len - real_len], [0, max_sen_len - real_len]], 'constant', constant_values=[-1, -1])
    
    #    根据postag的索引位置将postag的位置改为1
    for pt in postag:
        i, j = pt['loc'][0], pt['loc'][1]           #    取头尾指针
        m[i, j] = 1
        pass
    #    将m下三角部分置为-1，k=-1保留主对角线位置
    m[np.tril_indices(max_sen_len, k=-1)] = -1
    
    #    取所有1值的索引，并转为∑max_sen_len向量坐标
    loc_s, loc_e = np.where(m == 1)
    #    相对∑max_sen_len向量坐标 = 行坐标 * (max_sen_len + max_sen_len-1 +max_sen_len-2 + ... 当前行的前一行) + (列坐标 - 行坐标)
    idx_et2et = (max_sen_len + max_sen_len - loc_s + 1) * (loc_s) // 2 + (loc_e - loc_s)      

    #    取所有0值的索引，并转为∑max_sen_len向量坐标
    #    同样的，这里也不需要计算没握手的指针。算loss的时候用握手指针做成mask即可
    loc_s, loc_e = np.where(m == 0)
    #    相对∑max_sen_len向量坐标 = 行坐标 * (max_sen_len + max_sen_len-1 +max_sen_len-2 + ... 当前行的前一行) + (列坐标 - 行坐标)
    idx_no = (max_sen_len + max_sen_len - loc_s + 1) * (loc_s) // 2 + (loc_e - loc_s)
    #    根据正样本数量对负样本随机采样
    idx_no = np.random.choice(idx_no, len(idx_et2et))
    
    
    return idx_et2et, idx_no


#    将关系实体头尾握手信息转换为相对∑max_sen_len的向量索引信息
def sht2sht_vidx(text, spo_list, rel_size=len(rel.id_rel), max_sen_len=conf.TPLINKER.get_max_sentence_len()):
    '''将关系实体头尾握手信息转换为相对∑max_sen_len的向量索引信息
        @param text: 原文本字符串
        @param spo_list: 关系列表 [dict1, dict2, ...]
                            每个dict数据结构: {subject:{word:头实体字符串, loc:(头指针, 尾指针), type:类型},
                                              predicate:关系名称,
                                              object:{word:头实体字符串, loc:(头指针, 尾指针), type:类型}
        @param max_sen_len: 最大句子长度
        @return idx_sh2sh_1: narray(所有关系sh2sh==1的握手个数, 2)    2表示：[第几个关系, 握手信息在∑max_sen_len向量中的索引] 
                idx_sh2sh_2: 
                idx_st2st_1, 
                idx_st2st_2, 
                rel_no_val: 特殊的：若某一关系不存在握手数据，则该关系存储的索引为[关系id, -1]
                            没有记录存在关系，但该维度没有握手的索引。算loss的时候可以拿该维度的握手索引做成mask
                rel_has_1: 特殊的：所有存在握手1信息的维度列表
                rel_has_2: 特殊的：所有存在握手2信息的维度列表
    '''
    #    初始化len * len矩阵 [rel_size, 2, real_len, real_len]
    #    rel_size个关系，每个关系2个矩阵（sh2sh, st2st）
    real_len = len(text)
    m = np.zeros([rel_size, 2, real_len, real_len]) - 1
    #    取矩阵下三角部分，其他部分用0填充。取出来的就是下三角为-1，对角线和上三角为0
    m = np.tril(m, k=-1)
#     #    扩充为max_sen_len * max_sen_len矩阵
#     m = np.pad(m, [[0, 0], [0, 0], [0, max_sen_len - real_len], [0, max_sen_len - real_len]], 'constant', constant_values=[-1, -1])
    
    #    根据spo_list信息将对应位置改为1（subject在object前面）或者2（subject在object后面）
    rel_has_1 = set()                                   #    有握手1的维度
    rel_has_2 = set()                                   #    有握手2的维度
    rel_no_val = set(range(rel_size))                   #    没值的维度
    for spo in spo_list:
        s_loc = spo['subject']['loc']
        o_loc = spo['object']['loc']
        r_idx = rel.rel_id[spo['predicate']]
        
        #    如果subject的头指针比object小（subject在object前面，反映到m矩阵在上三角部分），则填充值为1，否则为2
        #    同时首尾指针位置交换（保证只填充上三角部分）
        val = 1
        if (s_loc[0] > o_loc[0]): 
            val = 2
            #    交换首尾指针
            t = s_loc
            s_loc = o_loc
            o_loc = t
            
            rel_has_2.add(r_idx)
            pass
        else:
            rel_has_1.add(r_idx)
            pass
        
        m[r_idx][0][s_loc[0], o_loc[0]] = val           #    头指针握手，设置sh2sh握手坐标（可能等于1，也可能等于2，但只有上三角部分会被设置）
        m[r_idx][1][s_loc[1], o_loc[1]] = val           #    尾指针握手，设置st2st握手坐标（可能等于1，也可能等于2，但只有上三角部分会被设置）
        
        if (r_idx in rel_no_val): rel_no_val.remove(r_idx)                        #    从没值的集合中删除该维度
        pass
    
    #    其实m矩阵已经可以直接拿去用了，但m太过稀疏，浪费空间太大。下面只保留他的索引信息
    
    #    将没值的维度置为-1，并且给没值的维度单独做成[维度, -1]的数据，用于表示该维度的不存在关系（节省存储空间）
    #    傻了，没有必要记录存在关系，但没有握手信息的索引数据。算loss的时可以用存在的握手索引数据做成mask，进一步节省空间
    m[list(rel_no_val)] = -1
    rel_no_val = np.array(list(rel_no_val))
    rel_has_1 = np.array(list(rel_has_1))
    rel_has_2 = np.array(list(rel_has_2))
#     rel_no_val = np.stack([rel_no_val, -np.ones_like(rel_no_val)], axis=-1)
    
    #    取所有握手和没握手的坐标
    #    取所有sh2sh==1的握手坐标，并转换为相对∑max_sen_len向量坐标
    idx_sh2sh_1 = np.stack(np.where(m[:, 0, :, :] == 1), axis=0).T        
    idx_sh2sh_1 = np.stack([idx_sh2sh_1[:, 0], (max_sen_len + max_sen_len - idx_sh2sh_1[:, 1] + 1) * (idx_sh2sh_1[:, 1]) // 2 + (idx_sh2sh_1[:, 2] - idx_sh2sh_1[:, 1])], axis=0).T
    #    取所有sh2sh==2的握手坐标（原下三角部分，被转上来的，实际场景中表示subject在object后面的情况）
    idx_sh2sh_2 = np.stack(np.where(m[:, 0, :, :] == 2), axis=0).T     
    idx_sh2sh_2 = np.stack([idx_sh2sh_2[:, 0], (max_sen_len + max_sen_len - idx_sh2sh_2[:, 1] + 1) * (idx_sh2sh_2[:, 1]) // 2 + (idx_sh2sh_2[:, 2] - idx_sh2sh_2[:, 1])], axis=0).T
    #    取所有sh2sh==0没有握手的位置坐标（这里取到的都是存在关系的维度）
    idx_sh2sh_no = np.stack(np.where(m[:, 0, :, :] == 0), axis=0).T
    idx_sh2sh_no = np.stack([idx_sh2sh_no[:, 0], (max_sen_len + max_sen_len - idx_sh2sh_no[:, 1] + 1) * (idx_sh2sh_no[:, 1]) // 2 + (idx_sh2sh_no[:, 2] - idx_sh2sh_no[:, 1])], axis=0).T
    #    根据idx_sh2sh_1 + idx_sh2sh_2的数量对idx_sh2sh_no随机采样
    idx_sh2sh_no = idx_sh2sh_no[np.random.choice(np.arange(len(idx_sh2sh_no)), len(idx_sh2sh_1) + len(idx_sh2sh_2))]
    
    #    取所有st2st==1的握手坐标
    idx_st2st_1 = np.stack(np.where(m[:, 1, :, :] == 1), axis=0).T
    idx_st2st_1 = np.stack([idx_st2st_1[:, 0], (max_sen_len + max_sen_len - idx_st2st_1[:, 1] + 1) * (idx_st2st_1[:, 1]) // 2 + (idx_st2st_1[:, 2] - idx_st2st_1[:, 1])], axis=0).T
    #    取所有st2st==2的握手坐标（原下三角部分，被转上来的，实际场景中表示subject在object后面的情况）
    idx_st2st_2 = np.stack(np.where(m[:, 1, :, :] == 2), axis=0).T     
    idx_st2st_2 = np.stack([idx_st2st_2[:, 0], (max_sen_len + max_sen_len - idx_st2st_2[:, 1] + 1) * (idx_st2st_2[:, 1]) // 2 + (idx_st2st_2[:, 2] - idx_st2st_2[:, 1])], axis=0).T
    #    取所有st2st==0没握手的位置坐标
    idx_st2st_no = np.stack(np.where(m[:, 1, :, :] == 0), axis=0).T
    idx_st2st_no = np.stack([idx_st2st_no[:, 0], (max_sen_len + max_sen_len - idx_st2st_no[:, 1] + 1) * (idx_st2st_no[:, 1]) // 2 + (idx_st2st_no[:, 2] - idx_st2st_no[:, 1])], axis=0).T
    #    根据idx_sh2sh_1 + idx_sh2sh_2的数量对idx_sh2sh_no随机采样
    idx_st2st_no = idx_st2st_no[np.random.choice(np.arange(len(idx_st2st_no)), len(idx_st2st_1) + len(idx_st2st_2))]
    
    return idx_sh2sh_1, idx_sh2sh_2, idx_sh2sh_no,\
            idx_st2st_1, idx_st2st_2, idx_st2st_no,\
            rel_no_val, rel_has_1, rel_has_2

text = 'sssss'
max_sen_len = 8
rel_size = 2
spo_list = [
                {
                    'subject':{'loc':(0,0)},
                    'predicate':'a',
                    'object':{'loc':(1,1)},
                },
                {
                    'subject':{'loc':(2,2)},
                    'predicate':'b',
                    'object':{'loc':(3,3)},
                },
                {
                    'subject':{'loc':(4,4)},
                    'predicate':'a',
                    'object':{'loc':(1,1)},
                },
            ]
rel.rel_id = {'a':0, 'b':1}
idx_sh2sh_1, idx_sh2sh_2, idx_sh2sh_no,\
            idx_st2st_1, idx_st2st_2, idx_st2st_no,\
            rel_no_val, rel_has_1, rel_has_2 = sht2sht_vidx(text, spo_list, rel_size, max_sen_len)

